#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import torch.nn as nn
from classy_vision.models import ClassyModel, register_model
import pytorchcv.model_provider as ptcv


@register_model("my_model")
class MyModel(ClassyModel):
    def __init__(self, arch, num_classes, **kwargs):
        super().__init__()

        model = ptcv.get_model(arch, **kwargs)
        self.features = model.features
        self.features.final_pool = nn.AdaptiveAvgPool2d((1, 1))
        # self.features.final_pool = nn.AdaptiveMaxPool2d((1, 1))
        in_features = model.output.in_features if hasattr(model.output, 'in_features') else model.output.in_channels
        self.output = nn.Linear(in_features=in_features, out_features=num_classes)

    def forward(self, x):
        x = self.features(x)
        x = x.view(int(x.size(0)), -1)
        x = self.output(x)
        return x

    @classmethod
    def from_config(cls, config):
        config.pop('name')
        arch = config.pop('arch')
        num_classes = config.pop('num_classes')
        return cls(arch, num_classes, **config)
